/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "DXHelper.h"
#include "BuddhabrotCS.h"
#include "DataBufferDX11.h"

/**
*/
BuddhabrotCS::BuddhabrotCS() : m_pCS( NULL ), m_pOutputUAV( NULL ), m_nInputWidth( 0 ), m_nInputHeight( 0 ), 
                               m_nOutputWidth( 0 ), m_nOutputHeight( 0 ), m_pConstantBuf( NULL ), m_bClearOutputUAV( true )
{
    assert( BUDDHABROT_SAMPLES_PER_FRAME == BUDDHABROT_CS_WIDTH*BUDDHABROT_CS_HEIGHT );
}

/**
*/
BuddhabrotCS::~BuddhabrotCS()
{
}

/**
*/
HRESULT BuddhabrotCS::OnCreateDevice( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;

    // compile compute shaders
    V_RETURN( CompileShaders( pd3dDevice ) );

    // Setup constant buffers
    D3D11_BUFFER_DESC Desc;
    Desc.Usage = D3D11_USAGE_DYNAMIC;
    Desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    Desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    Desc.MiscFlags = 0;    
    Desc.ByteWidth = sizeof( ConstantBuf );
    V_RETURN( pd3dDevice->CreateBuffer( &Desc, NULL, &m_pConstantBuf ) );

    return hr;
}

/**
*/
HRESULT BuddhabrotCS::CompileShaders( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    ID3DBlob* pBlob = NULL;

    // Compile the Buddhabrot Compute Shader
    V_RETURN( CompileShaderFromFile( L"BuddhabrotCS.hlsl", "BuddhabrotCS", "cs_5_0", &pBlob ) );
    
    // Create the Buddhabrot Compute Shader
    V_RETURN( pd3dDevice->CreateComputeShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pCS ) );

    SAFE_RELEASE( pBlob );

    return hr;
}

/**
*/
HRESULT BuddhabrotCS::OnResizedSwapChain( ID3D11Device* pd3dDevice, DataBufferDX11* pDataBuffer )
{
    HRESULT hr = S_OK;

    // destroy old uav
    SAFE_RELEASE( m_pOutputUAV );

    // destory old random number buffer
    m_randomNumBuffer.OnDestroyDevice();

    // save output width, height
    m_nOutputWidth = pDataBuffer->GetWidth();
    m_nOutputHeight = pDataBuffer->GetHeight();

    m_nInputWidth = BUDDHABROT_CS_WIDTH;
    m_nInputHeight = BUDDHABROT_CS_HEIGHT;

    // create uav 
    V_RETURN( pd3dDevice->CreateUnorderedAccessView( pDataBuffer->GetpBufferGPU(), NULL, &m_pOutputUAV ) );

    // create random number data for buddhabrot sampling
    V_RETURN( m_randomNumBuffer.OnResizedSwapChain( pd3dDevice, m_nInputWidth, m_nInputHeight, m_cb.fMinRe, m_cb.fMaxRe, m_cb.fMinIm, m_cb.fMaxIm ) );

    m_bClearOutputUAV = true;

    return hr;
}

/**
*/
void BuddhabrotCS::OnFrameRender( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext,
                                  double fTime, float fElapsedTime, DataBufferDX11* pDataBuffer )
{
    HRESULT hr = S_OK;

    UpdateConstantBuf( m_nInputWidth, m_nInputHeight, m_nOutputWidth, m_nOutputHeight );

    // clear the output uav when the user translates or scales the Buddhrabrot.
    if( m_bClearOutputUAV ) {
        const UINT clearValues[ 4 ] = { 0, 0, 0, 0 };
        pd3dImmediateContext->ClearUnorderedAccessViewUint( m_pOutputUAV, clearValues );
        m_bClearOutputUAV = false;
        m_randomNumBuffer.UpdateBuffer( pd3dDevice, pd3dImmediateContext, m_cb.fMinRe, m_cb.fMaxRe, m_cb.fMinIm, m_cb.fMaxIm );
    }

    m_randomNumBuffer.UpdateBuffer( pd3dDevice, pd3dImmediateContext, m_cb.fMinRe, m_cb.fMaxRe, m_cb.fMinIm, m_cb.fMaxIm );

    // set constant buffer
    D3D11_MAPPED_SUBRESOURCE MappedResource;            
    V( pd3dImmediateContext->Map( m_pConstantBuf, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) );
    ConstantBuf* buf = (ConstantBuf*)MappedResource.pData;
    *buf = m_cb;
    pd3dImmediateContext->Unmap( m_pConstantBuf, 0 );
    ID3D11Buffer* ppCB[1] = { m_pConstantBuf };
    pd3dImmediateContext->CSSetConstantBuffers( 0, 1, ppCB );

    pd3dImmediateContext->CSSetShader( m_pCS, NULL, 0 );

    // compute shader input
    ID3D11ShaderResourceView* aRViews[ 1 ] = { m_randomNumBuffer.GetpSRV() };
    pd3dImmediateContext->CSSetShaderResources( 0, 1, aRViews );

    // compute shader output
    ID3D11UnorderedAccessView* aUAViews[ 1 ] = { m_pOutputUAV };
    pd3dImmediateContext->CSSetUnorderedAccessViews( 0, 1, aUAViews, (UINT*)(&aUAViews) );

    // run compute shader
    pd3dImmediateContext->Dispatch( (m_nInputWidth)/10, (m_nInputHeight)/10, 1 );

    // clear uav
    ID3D11UnorderedAccessView* aUAViewsClear[ 1 ] = { NULL };
    pd3dImmediateContext->CSSetUnorderedAccessViews( 0, 1, aUAViewsClear, (UINT*)(&aUAViewsClear) );

    // debug code. wanted to see content of the buddhabrot output buffer
#if 0
    {
        ID3D11Buffer* pDebugBuf = CreateAndCopyToDebugBuf( pd3dDevice, pd3dImmediateContext, pDataBuffer->GetpBufferGPU() );
        D3D11_MAPPED_SUBRESOURCE MappedResource;     
        V( pd3dImmediateContext->Map( pDebugBuf, 0, D3D11_MAP_READ, 0, &MappedResource ) );
        BufType* pData = reinterpret_cast< BufType* >( MappedResource.pData );
        for( int i=0; i<pDataBuffer->GetWidth()*pDataBuffer->GetHeight(); ++i ) 
        {
            BufType bufType = pData[ i ];
            if( bufType.v[ 0 ] > 10000 ) {
                int nBreak = 0;
            }
        }
        pd3dImmediateContext->Unmap( pDebugBuf, 0 );
        SAFE_RELEASE( pDebugBuf ); 
    }
#endif
}

/**
*/
void BuddhabrotCS::OnDestroyDevice()
{
    SAFE_RELEASE( m_pCS );
    SAFE_RELEASE( m_pOutputUAV );
    SAFE_RELEASE( m_pConstantBuf );
    m_randomNumBuffer.OnDestroyDevice();
}
